#pragma once

#include <nanobdd/Common.h>
#include <nanobdd/Node.h>
#include <tbb/concurrent_vector.h>
#include <execution>
#include <iostream>
#include <map>
#include <shared_mutex>
#include <unordered_map>
#include <vector>

namespace nanobdd {

class Bucket {
 public:
  Bucket(size_t size)
      : size_(size), currSize_(0), mutex_(std::make_unique<std::shared_mutex>()) {
    for (int i = 0; i < 34; i++) {
      levelNodeMap_.push_back({});
    }
  }

  // TBB approach
  Node*
  operator()(int level, Node* low, Node* high) {
    {
    std::shared_lock<std::shared_mutex> lock(*mutex_);
    for (auto& node : nodes_) {
      if (node.low == low && node.high == high && node.level == level) {
        return &node;
      }
    }
    }
    std::lock_guard<std::shared_mutex> lock(*mutex_);
    for (auto& node : nodes_) {
      if (node.low == low && node.high == high && node.level == level) {
        return &node;
      }
    }
    Node node;
    node.level = level;
    node.low = low;
    node.high = high;
    auto it = nodes_.push_back(std::move(node));
    return &(*it);
  }

  // Node*
  // operator()(int level, Node* low, Node* high) {
  //   // search in nodes
  //   {
  //     std::shared_lock<std::shared_mutex> lock(*mutex_);
  //     for (int i = 0; i < currSize_; ++i) {
  //       auto& node = nodes_.at(i);
  //       if (node.low == low && node.high == high && node.level == level) {
  //         return &node;
  //       }
  //     }
  //   }

  //   std::lock_guard<std::shared_mutex> lock(*mutex_);
  //   // search again
  //   for (int i = 0; i < currSize_; ++i) {
  //     auto& node = nodes_.at(i);
  //     if (node.low == low && node.high == high && node.level == level) {
  //       return &node;
  //     }
  //   }

  //   // create the node
  //   auto& node = nodes_.at(currSize_);
  //   node.level = level;
  //   node.low = low;
  //   node.high = high;
  //   currSize_++;
  //   if (currSize_ > size_) {
  //     std::cout << "bucket full" << std::endl;
  //   }

  //   return &node;
  // };

  // Node*
  // operator()(int level, Node* low, Node* high) {
  //   auto& map = levelNodeMap_.at(level);
  //   {
  //     std::shared_lock<std::shared_mutex> lock(*mutex_);
  //     auto it1 = map.find(low);
  //     if (it1 != map.end()) {
  //       auto it2 = it1->second.find(high);
  //       if (it2 != it1->second.end()) {
  //         return &(it2->second);
  //       }
  //     }
  //   }
  //   std::lock_guard<std::shared_mutex> lock(*mutex_);
  //   auto it1 = map.find(low);
  //   if (it1 != map.end()) {
  //     auto it2 = it1->second.find(high);
  //     if (it2 != it1->second.end()) {
  //       return &(it2->second);
  //     }
  //   }
  //   auto& node = levelNodeMap_[level][low][high];
  //   node.level_ = level;
  //   node.low_ = low;
  //   node.high_ = high;

  //   return &node;
  // }

 private:
  size_t size_;
  size_t currSize_; // current occupied number of nodes
  // std::vector<Node> nodes_;
  tbb::concurrent_vector<Node> nodes_;
  // std::unique_ptr<std::mutex> mutex_;
  std::unique_ptr<std::shared_mutex> mutex_;
  std::map<int, std::map<Node*, std::map<Node*, Node>>> mapNodes_;
  std::vector<std::unordered_map<Node*, std::unordered_map<Node*, Node>>>
      levelNodeMap_;
};

class NodeTable {
 public:
  NodeTable(size_t tableSize, int bucketSize) : tableSize_(tableSize) {
    for (int i = 0; i < tableSize; ++i) {
      buckets_.emplace_back(bucketSize);
    }
  };

  /**
   * Get or create a node
   */
  Node*
  operator()(uint32_t level, Node* low, Node* high) {
    auto hash = HASH_O_3(
                    level,
                    reinterpret_cast<uintptr_t>(low),
                    reinterpret_cast<uintptr_t>(high)) %
        tableSize_;
    auto& bucket = buckets_.at(hash);

    return bucket(level, low, high);
  };

 private:
  size_t tableSize_;
  std::vector<Bucket> buckets_;
};

} // namespace nanobdd